package tcp

import (
	"context"
	"github.com/DiracLee/dires-go/logger"
	"github.com/DiracLee/dires-go/tcp/handler"
	"net"
	"os"
	"os/signal"
	"sync"
	"sync/atomic"
	"syscall"
)

type Listener struct {
	net.Listener
	closed uint32
}

func (l *Listener) Close() error {
	if !atomic.CompareAndSwapUint32(&l.closed, 0, 1) {
		// avoid closing a closed listener
		return nil
	}
	return l.Listener.Close()
}

func ListenAndServerUntilExitSignal(ctx context.Context, config *Config, h handler.Handler) error {
	if h == nil {
		h = handler.NewEchoHandler()
	}
	closeChan := make(chan struct{})
	signalChan := make(chan os.Signal)
	signal.Notify(signalChan, syscall.SIGHUP, syscall.SIGQUIT, syscall.SIGTERM, syscall.SIGINT)
	go func() {
		select {
		case <-signalChan:
			closeChan <- struct{}{}
		}
	}()
	listener, err := net.Listen("tcp", config.Address)
	if err != nil {
		logger.Errorf("[ListenAndServerUntilExitSignal] TCP fail to listen at address(%s), err: %v", config.Address, err)
		return err
	}
	logger.Infof("[ListenAndServerUntilExitSignal] server bind in %s, start listening ...", config.Address)
	ListenAndServe(ctx, &Listener{
		Listener: listener,
		closed:   0,
	}, h, closeChan)
	return nil
}

func ListenAndServe(ctx context.Context, listener *Listener, h handler.Handler, closeChan chan struct{}) {
	defer func() {
		_ = listener.Close()
		_ = h.Close()
	}()

	go func() {
		select {
		case <-closeChan:
			_ = listener.Close()
			_ = h.Close()
		}
	}()

	wg := sync.WaitGroup{}
	for {
		conn, err := listener.Accept()
		if err != nil {
			logger.Errorf("[ListenAndServe] listener failed to Accept(), err: %v", err)
			break
		}
		logger.Infof("[ListenAndServe] accepted a connection (%v)", conn.RemoteAddr().String())
		wg.Add(1)
		go func() {
			defer func() {
				wg.Done()
			}()
			h.Handle(ctx, conn)
		}()
	}
}
